import random
import torch
import torch.nn.functional as F

def random_crop(seqs: torch.Tensor, crop_size):
    """
    seqs: (T*C, H, W)
    """
    TC, H, W = seqs.size()

    left = random.randint(0, W - crop_size[0])
    top = random.randint(0, H - crop_size[1])
    right = left + crop_size[0]
    bottom = top + crop_size[1]

    return seqs[:, top: bottom, left: right]


def pad(x, p = 64):
    h, w = x.size(-2), x.size(-1)

    new_h = (h + p - 1) // p * p     
    new_w = (w + p - 1) // p * p

    padding_left = (new_w - w) // 2
    padding_right = new_w - w - padding_left     
    padding_top = (new_h - h) // 2     
    padding_bottom = new_h - h - padding_top     
    padding = (padding_left, padding_right, padding_top, padding_bottom)     

    x = F.pad(
        x,         
        padding,         
        mode="constant",         
        value=0,     
    )
    return x, padding


def crop(x, padding):
    return F.pad(x, tuple(-p for p in padding))


def random_flip(seqs: torch.Tensor):
    """
    seqs: (T*C, H, W)
    """
        
    # horizontal
    if random.randint(0, 1) == 1:
        seqs = torch.flip(seqs, [1])
    
    # vertical
    if random.randint(0, 1) == 1:
        seqs = torch.flip(seqs, [2])

    return seqs
